#include <ros/ros.h>
#include <moveit/move_group_interface/move_group_interface.h>
#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include <moveit_msgs/CollisionObject.h>
#include <moveit_msgs/AttachedCollisionObject.h>
#include <geometry_msgs/Pose.h>
#include <vector>

void addObstacles(moveit::planning_interface::PlanningSceneInterface& planning_scene_interface) {
    // Define obstacle A
    moveit_msgs::CollisionObject obstacle_a;
    obstacle_a.id = "obstacle_a";
    obstacle_a.header.frame_id = "base_link";

    shape_msgs::SolidPrimitive box_primitive;
    box_primitive.type = shape_msgs::SolidPrimitive::BOX;
    box_primitive.dimensions = {0.3, 0.3, 0.05};

    geometry_msgs::Pose box_pose;
    box_pose.position.x = 0.3;
    box_pose.position.y = 0.3;
    box_pose.position.z = 0.2;
    box_pose.orientation.w = 1.0;

    obstacle_a.primitives.push_back(box_primitive);
    obstacle_a.primitive_poses.push_back(box_pose);
    obstacle_a.operation = moveit_msgs::CollisionObject::ADD;

    // Define obstacle B
    moveit_msgs::CollisionObject obstacle_b = obstacle_a;
    obstacle_b.id = "obstacle_b";
    obstacle_b.primitive_poses[0].position.x = 0.3;
    obstacle_b.primitive_poses[0].position.y = -0.2;

    // Define obstacle C
    moveit_msgs::CollisionObject obstacle_c = obstacle_a;
    obstacle_c.id = "obstacle_c";
    obstacle_c.primitive_poses[0].position.x = 0.3;
    obstacle_c.primitive_poses[0].position.y = 0.2;

    // Add obstacles to the planning scene
    // planning_scene_interface.applyCollisionObjects({obstacle_a, obstacle_b, obstacle_c});
    planning_scene_interface.applyCollisionObjects({obstacle_a});
}

void moveToPosition(moveit::planning_interface::MoveGroupInterface& move_group, const std::vector<double>& position) {
    geometry_msgs::Pose target_pose = move_group.getCurrentPose().pose;
    target_pose.position.x = position[0];
    target_pose.position.y = position[1];
    target_pose.position.z = position[2];

    move_group.setPoseTarget(target_pose);

    moveit::planning_interface::MoveGroupInterface::Plan plan;
    if (move_group.plan(plan) == moveit::planning_interface::MoveItErrorCode::SUCCESS) {
        move_group.execute(plan);
    } else {
        ROS_WARN("Failed to plan to target position");
    }

    move_group.clearPoseTargets();
}

void attachObject(moveit::planning_interface::PlanningSceneInterface& planning_scene_interface,
                  moveit::planning_interface::MoveGroupInterface& move_group,
                  const std::string& object_id) {
    moveit_msgs::AttachedCollisionObject attached_object;
    attached_object.object.header.frame_id = "tool0";
    attached_object.link_name ="tool0";
    attached_object.object.id = object_id;
    attached_object.object.operation = moveit_msgs::CollisionObject::ADD;

    attached_object.object.primitives.resize(1);
    attached_object.object.primitive_poses.resize(1);
    attached_object.object.primitives[0].type = shape_msgs::SolidPrimitive::BOX;
    attached_object.object.primitives[0].dimensions = {0.05, 0.05, 0.05};
    attached_object.object.primitive_poses[0].position.x = 0.0;
    attached_object.object.primitive_poses[0].position.y = -0.05;
    attached_object.object.primitive_poses[0].position.z = 0.0;
    planning_scene_interface.applyAttachedCollisionObject(attached_object);
}

void detachObject(moveit::planning_interface::PlanningSceneInterface& planning_scene_interface,
                  const std::string& object_id) {
    planning_scene_interface.removeCollisionObjects(std::vector<std::string>{object_id});
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "robot_grasp_demo_cpp");
    ros::NodeHandle nh;
    ros::AsyncSpinner spinner(1);
    spinner.start();

    moveit::planning_interface::MoveGroupInterface move_group("manipulator");
    moveit::planning_interface::PlanningSceneInterface planning_scene_interface;
    move_group.setMaxAccelerationScalingFactor(0.8);
    move_group.setMaxVelocityScalingFactor(0.8);

    ros::Duration(2.0).sleep();  // Ensure the planning scene is ready

    // Add obstacles to the scene
    addObstacles(planning_scene_interface);
    ros::Duration(2.0).sleep();

    // 1. Move to position p1
    std::vector<double> p1 = {0.3, 0.3, 0.3};
    moveToPosition(move_group, p1);

    // 2. Attach obstacle A
    attachObject(planning_scene_interface, move_group, "obstacle_d");
    ros::Duration(2.0).sleep();

    // 3. Move to position p2
    std::vector<double> p2 = {-0.3, -0.3, 0.3};
    moveToPosition(move_group, p2);

    // 4. Detach obstacle A
    detachObject(planning_scene_interface, "obstacle_d");
    ros::Duration(2.0).sleep();

    // 5. Return to home position
    std::vector<double> home = {0.3, 0.3, 0.3};
    moveToPosition(move_group, home);

    ROS_INFO("Task Completed!");

    return 0;
}
